import pandas as pd
import matplotlib.pyplot as plt
import os
import numpy as np
from scipy.signal import savgol_filter
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from pyts.approximation import PiecewiseAggregateApproximation
from matplotlib import rcParams
from data_loader import DataLoader

# 配置matplotlib以支持中文和负号
config = {
    "font.family": 'serif',
    "font.size": 18,
    "mathtext.fontset": 'stix',
    "font.serif": ['SimSun'],  # 宋体
    'axes.unicode_minus': False  # 处理负号
}
try:
    rcParams.update(config)
except Exception as e:
    print(f"字体配置失败: {e}")


class SpectraPreprocessor:
    """
    九种近红外光谱预处理方法
    msc_optimized， snv， sg， D1， D2，mean_centralization，
    min_max_normalization， move_avg， paa
    """

    @staticmethod
    # 多元散射校正 (MSC)
    def msc(data):
        # 检查输入数据是否为DataFrame类型，如果是则转换为NumPy数组
        if isinstance(data, pd.DataFrame):
            data = data.values
        # 定义一个字典，用于存储处理后的数据
        n_samples, n_features = data.shape
        # 计算参考光谱
        mean_spectrum = np.mean(data, axis=0)
        # 计算参考光谱的斜率和截距
        k = np.zeros(n_samples)
        b = np.zeros(n_samples)
        corrected_spectra = np.zeros_like(data)

        model = LinearRegression()
        # 遍历每个样本，进行多元散射校正
        for i in range(n_samples):
            model.fit(mean_spectrum.reshape(-1, 1), data[i, :].reshape(-1, 1))
            k[i] = model.coef_
            b[i] = model.intercept_
            # 使用广播直接校正
            corrected_spectra[i, :] = (data[i, :] - b[i]) / k[i]

        return corrected_spectra

    @staticmethod
    # 标准正态变换 (SNV)
    def snv(data):
        # 检查输入数据是否为DataFrame类型，如果是则转换为NumPy数组
        if isinstance(data, pd.DataFrame):
            data = data.values
        return (data - np.mean(data, axis=1)[:, np.newaxis]) / np.std(data, axis=1)[:, np.newaxis]

    @staticmethod
    # Savitzky-Golay 滤波 (SG)
    def sg(data, window_length=13, polyorder=3):
        if isinstance(data, pd.DataFrame):
            data = data.values
        signal_f = savgol_filter(data, window_length, polyorder)
        return signal_f

    @staticmethod
    # 一阶差分算法 (一阶导数+sg滤波)
    def D1(data):
        if isinstance(data, pd.DataFrame):
            data = data.values
        return savgol_filter(data, window_length=13, polyorder=3, deriv=1)  # deriv=1表示二阶导数

    @staticmethod
    # 二阶差分算法 (导数+sg滤波)
    def D2(data):
        if isinstance(data, pd.DataFrame):
            data = data.values
        return savgol_filter(data, window_length=13, polyorder=3, deriv=2)  # deriv=2表示二阶导数

    @staticmethod
    # 均值中心化（将每列数据减去其均值，使数据均值为0）
    def mean_centralization(data):
        # 检查输入数据是否为DataFrame类型，如果是则转换为NumPy数组
        if isinstance(data, pd.DataFrame):
            data = data.values
        # 计算列均值
        column_mean = np.mean(data, axis=0)
        # 通过广播实现均值中心化
        centered_data = data - column_mean
        return centered_data

    @staticmethod
    # 最大最小归一化
    def min_max_normalization(data):
        if isinstance(data, pd.DataFrame):
            data = data.values
        min_vals = np.min(data, axis=1, keepdims=True)
        max_vals = np.max(data, axis=1, keepdims=True)
        ranges = max_vals - min_vals
        # 避免除零错误，对于常数行直接设置为零
        ranges[ranges == 0] = 1
        return (data - min_vals) / ranges

    @staticmethod
    # 滑动平均滤波
    def move_avg(data, n=5, mode="valid"):
        if isinstance(data, pd.DataFrame):
            data = data.values
        from scipy.signal import convolve2d
        window = np.ones((1, n)) / n  # 定义窗口 (1D 平均窗口)
        tmp = convolve2d(data, window, mode=mode)
        if tmp.shape[1] < data.shape[1]:
            pad_length = data.shape[1] - tmp.shape[1]  # 计算长度差
            tmp = np.pad(tmp, ((0, 0), (0, pad_length)), mode='edge')  # 用最后一个值填补
        return tmp

    @staticmethod
    # 分段聚合近似法 (PAA)降维, 降维后wave会与之不匹配，需要重新取索引
    def paa(data, wave, window_size=2):
        # 处理data
        paa = PiecewiseAggregateApproximation(window_size=window_size)
        paa_data = paa.transform(data)
        # 处理wave
        indices = np.linspace(0, len(wave) - 1, num=paa_data.shape[1], dtype=int)
        matched_wave = wave[indices]  # 从原波长中取对应索引的数据
        return paa_data, matched_wave

    @staticmethod
    def plot_spectra(wave, data, title="Spectra", xlabel="Wavelength (nm)", ylabel="Reabsorbance"):
        """Plot spectral data."""
        plt.figure(figsize=(15, 10))
        plt.plot(wave, data.T, alpha=0.5)
        plt.title(title, fontsize=16)
        plt.xlabel(xlabel, fontsize=12)
        plt.ylabel(ylabel, fontsize=12)
        plt.grid(True)
        plt.show()


if __name__ == '__main__':
    # 加载鸭梨数据, 返回值分别为对应的光谱、标签、波长
    spectra, label, wave = DataLoader.load_peach_spectra()
    pre = SpectraPreprocessor()
    # # 绘制原始光谱图
    # pre.plot_spectra(wave, spectra, title="原始光谱图")
    # # 多元散射校正 (MSC)
    # msc_spectra = pre.msc(spectra)
    # pre.plot_spectra(wave, msc_spectra, title="MSC Processed Spectra")
    # # 标准正态变换 (SNV)
    # snv_spectra = pre.snv(spectra)
    # pre.plot_spectra(wave, spectra, title="snv Spectra")
    # # Savitzky-Golay 滤波 (SG)
    sg_spectar = pre.sg(spectra)

    pre.plot_spectra(wave, sg_spectar, title="sg Spectra")
    # # 一阶差分算法 (一阶导数+sg滤波)
    # d1_spectar = pre.D1(spectra)
    # pre.plot_spectra(wave, d1_spectar, title="D1 Spectra")
    # # 二阶差分算法 (导数+sg滤波)
    # d2_spectar = pre.D2(spectra)
    # pre.plot_spectra(wave, d2_spectar, title="D2 Spectra")
    # # 均值中心化
    # mean_spectra = pre.mean_centralization(spectra)
    # pre.plot_spectra(wave, mean_spectra, title="mean Spectra")
    # # 最大最小归一化
    # min_max_spectra = pre.min_max_normalization(spectra)
    # pre.plot_spectra(wave, min_max_spectra, title="min_max Spectra")
    # # 滑动平均滤波
    # move_spectra = pre.move_avg(spectra, n=5)
    # pre.plot_spectra(wave, move_spectra, title="move Spectra")
    # 分段聚合近似法 (PAA)降维
    # paa_spectra, matched_wave = pre.paa(spectra, wave, window_size=2)
    # pre.plot_spectra(matched_wave, paa_spectra, title="paa Spectra")
    # snv_spectra = pre.snv(spectra)
    # sg_spectra = pre.sg(snv_spectra)
    # pre.plot_spectra(wave, sg_spectra, title="precprocess Spectra")











